from typing import Dict, List, Optional, Set
from chernc.treediff.gumtree.actions.edit_script import EditScript
from chernc.treediff.gumtree.actions.edit_script_generator import EditScriptGenerator
from chernc.treediff.gumtree.actions.model.delete import Delete
from chernc.treediff.gumtree.actions.model.insert import Insert
from chernc.treediff.gumtree.actions.model.move import Move
from chernc.treediff.gumtree.actions.model.update import Update
from chernc.treediff.gumtree.matchers.mapping import Mapping
from chernc.treediff.gumtree.matchers.mapping_store import MappingStore
from chernc.treediff.gumtree.tree.fake_tree import FakeTree
from chernc.treediff.gumtree.tree.tree import TreeLike
from chernc.treediff.gumtree.tree.tree_utils import TreeUtils


class ChawatheScriptGenerator(EditScriptGenerator):
    def __init__(self):
        self.orig_src: Optional[TreeLike] = None
        self.cpy_src: Optional[TreeLike] = None
        self.orig_dst: Optional[TreeLike] = None
        self.orig_mappings: Optional[MappingStore[TreeLike]] = None
        self.cpy_mappings: Optional[MappingStore[TreeLike]] = None
        self.dst_in_order: Optional[Set[TreeLike]] = None
        self.src_in_order: Optional[Set[TreeLike]] = None
        self.actions: Optional[EditScript] = None
        self.orig_to_copy: Optional[Dict[TreeLike, TreeLike]] = None
        self.copy_to_orig: Optional[Dict[TreeLike, TreeLike]] = None

    def compute_actions(self, mappings: 'MappingStore') -> EditScript:
        self.init_with(mappings)
        self.generate()
        if self.actions is None:
            raise Exception("The actions result is None.")
        return self.actions

    def init_with(self, ms: MappingStore[TreeLike]):
        self.orig_src = ms.src
        self.cpy_src = self.orig_src.deep_copy()
        self.orig_dst = ms.dst
        self.orig_mappings = ms

        self.orig_to_copy = {}
        self.copy_to_orig = {}
        cpy_tree_iterator = TreeUtils.pre_order_iterator(self.cpy_src)
        for orig_tree in TreeUtils.pre_order(self.orig_src):
            cpy_tree = next(cpy_tree_iterator)
            self.orig_to_copy[orig_tree] = cpy_tree
            self.copy_to_orig[cpy_tree] = orig_tree

        self.cpy_mappings = MappingStore(ms.src, ms.dst)
        for m in self.orig_mappings:
            self.cpy_mappings.add_mapping(self.orig_to_copy[m.src], m.dst)

    def generate(self):
        if self.cpy_src is None:
            raise Exception("The cpy_src is None. Maybe ChawatheScriptGenerator has not been initialized.")
        if self.orig_dst is None:
            raise Exception("The orig_dst is None. Maybe ChawatheScriptGenerator has not been initialized.")
        if self.cpy_mappings is None:
            raise Exception("The cpy_mappings is None. Maybe ChawatheScriptGenerator has not been initialized.")
        if self.copy_to_orig is None:
            raise Exception("The copy_to_orig is None. Maybe ChawatheScriptGenerator has not been initialized.")
        src_fake_root = FakeTree(self.cpy_src)
        dst_fake_root = FakeTree(self.orig_dst)
        self.cpy_src.set_parent(src_fake_root)
        self.orig_dst.set_parent(dst_fake_root)

        self.actions = EditScript()
        self.dst_in_order = set()
        self.src_in_order = set()

        self.cpy_mappings.add_mapping(src_fake_root, dst_fake_root)

        bfs_dst = TreeUtils.breadth_first(self.orig_dst)
        for x in bfs_dst:
            w = None
            y = x.get_parent()
            z = self.cpy_mappings.get_src_for_dst(y)

            if not self.cpy_mappings.is_dst_mapped(x):
                k = self.find_pos(x)
                # Insertion case: insert new node
                w = FakeTree()
                # Furnish x instead of w to use real nodes from the second tree
                ins = Insert(x, self.copy_to_orig.get(z), k)
                self.actions.add(ins)
                self.copy_to_orig[w] = x
                self.cpy_mappings.add_mapping(w, x)
                z.insert_child(w, k)
            else:
                w = self.cpy_mappings.get_src_for_dst(x)
                if x != self.orig_dst:  # Handle root case separately
                    v = w.get_parent()
                    if w.get_label() != x.get_label():
                        self.actions.add(Update(self.copy_to_orig.get(w), x.get_label()))
                        w.set_label(x.get_label())
                    if z != v:
                        k = self.find_pos(x)
                        mv = Move(self.copy_to_orig.get(w), self.copy_to_orig.get(z), k)
                        self.actions.add(mv)
                        old_k = w.position_in_parent()
                        w.get_parent().get_children().pop(old_k)
                        z.insert_child(w, k)

            self.src_in_order.add(w)
            self.dst_in_order.add(x)
            self.align_children(w, x)

        for w in self.cpy_src.post_order():
            if not self.cpy_mappings.is_src_mapped(w):
                self.actions.add(Delete(self.copy_to_orig.get(w)))

    def align_children(self, w: TreeLike, x: TreeLike):
        if self.src_in_order is None:
            raise Exception("The src_in_order is None. Maybe ChawatheScriptGenerator has not been initialized.")
        if self.dst_in_order is None:
            raise Exception("The dst_in_order is None. Maybe ChawatheScriptGenerator has not been initialized.")
        if self.cpy_mappings is None:
            raise Exception("The cpy_mappings is None. Maybe ChawatheScriptGenerator has not been initialized.")
        if self.copy_to_orig is None:
            raise Exception("The copy_to_orig is None. Maybe ChawatheScriptGenerator has not been initialized.")
        if self.actions is None:
            raise Exception("The copy_to_orig is None. Maybe ChawatheScriptGenerator has not been initialized.")
        self.src_in_order.difference_update(w.get_children())
        self.dst_in_order.difference_update(x.get_children())

        s1 = [c for c in w.get_children() if self.cpy_mappings.is_src_mapped(c) and self.cpy_mappings.get_dst_for_src(c) in x.get_children()]
        s2 = [c for c in x.get_children() if self.cpy_mappings.is_dst_mapped(c) and self.cpy_mappings.get_src_for_dst(c) in w.get_children()]

        lcs = self.lcs(s1, s2)

        for m in lcs:
            self.src_in_order.add(m.src)
            self.dst_in_order.add(m.dst)

        for b in s2:
            for a in s1:
                if self.cpy_mappings.has(a, b):
                    if Mapping(a, b) not in lcs:
                        a.get_parent().get_children().remove(a)
                        k = self.find_pos(b)
                        mv = Move(self.copy_to_orig.get(a), self.copy_to_orig.get(w), k)
                        self.actions.add(mv)
                        w.get_children().insert(k, a)
                        a.set_parent(w)
                        self.src_in_order.add(a)
                        self.dst_in_order.add(b)

    def find_pos(self, x: 'TreeLike') -> int:
        y = x.get_parent()
        siblings = y.get_children()

        for c in siblings:
            if c in self.dst_in_order:
                if c == x:
                    return 0
                else:
                    break

        xpos = x.position_in_parent()
        v = None
        for i in range(xpos):
            c = siblings[i]
            if c in self.dst_in_order:
                v = c

        if v is None:
            return 0

        u = self.cpy_mappings.get_src_for_dst(v)
        upos = u.position_in_parent()
        return upos + 1

    def lcs(self, x: List[TreeLike], y: List[TreeLike]) -> List[Mapping[TreeLike]]:
        if self.cpy_mappings is None:
            raise Exception("The cpy_mappings is None. Maybe ChawatheScriptGenerator has not been initialized.")
        m = len(x)
        n = len(y)
        lcs = []

        opt = [[0] * (n + 1) for _ in range(m + 1)]
        for i in range(m - 1, -1, -1):
            for j in range(n - 1, -1, -1):
                if self.cpy_mappings.get_src_for_dst(y[j]) == x[i]:
                    opt[i][j] = opt[i + 1][j + 1] + 1
                else:
                    opt[i][j] = max(opt[i + 1][j], opt[i][j + 1])

        i, j = 0, 0
        while i < m and j < n:
            if self.cpy_mappings.get_src_for_dst(y[j]) == x[i]:
                lcs.append(Mapping(x[i], y[j]))
                i += 1
                j += 1
            elif opt[i + 1][j] >= opt[i][j + 1]:
                i += 1
            else:
                j += 1

        return lcs
